import matplotlib.pyplot as plt
import pandas as pd
import numpy as np
import os
from typing import Dict, List, Any, Optional, Tuple
import logging

class Visualizer:
    """回测结果可视化工具"""
    
    def __init__(self, output_dir: str = "./output"):
        """
        初始化可视化工具
        
        参数:
            output_dir: 图表输出目录
        """
        self.output_dir = output_dir
        self.logger = logging.getLogger(__name__)
        
        # 创建输出目录
        os.makedirs(output_dir, exist_ok=True)
        
        # 设置绘图样式
        plt.style.use('seaborn-v0_8-darkgrid')
        
    def plot_equity_curve(self, backtest_data: pd.DataFrame, benchmark_data: Optional[pd.DataFrame] = None,
                         title: str = "权益曲线", filename: str = "equity_curve.png") -> None:
        """
        绘制权益曲线
        
        参数:
            backtest_data: 回测数据，包含cum_returns列
            benchmark_data: 基准数据，包含cum_returns列
            title: 图表标题
            filename: 输出文件名
        """
        try:
            plt.figure(figsize=(12, 6))
            
            # 绘制策略权益曲线
            if 'cum_returns' in backtest_data.columns:
                plt.plot(backtest_data.index, backtest_data['cum_returns'], 
                        label='策略', linewidth=2)
            else:
                # 计算累积收益
                backtest_data['cum_returns'] = (1 + backtest_data['returns']).cumprod()
                plt.plot(backtest_data.index, backtest_data['cum_returns'], 
                        label='策略', linewidth=2)
            
            # 绘制基准权益曲线
            if benchmark_data is not None:
                if 'cum_returns' in benchmark_data.columns:
                    plt.plot(benchmark_data.index, benchmark_data['cum_returns'], 
                            label='基准', linewidth=2, alpha=0.7)
                else:
                    # 计算基准累积收益
                    benchmark_data['cum_returns'] = (1 + benchmark_data['returns']).cumprod()
                    plt.plot(benchmark_data.index, benchmark_data['cum_returns'], 
                            label='基准', linewidth=2, alpha=0.7)
            
            plt.title(title, fontsize=14)
            plt.xlabel('日期', fontsize=12)
            plt.ylabel('累积收益', fontsize=12)
            plt.legend()
            plt.grid(True)
            
            # 保存图表
            plt.savefig(os.path.join(self.output_dir, filename), dpi=300, bbox_inches='tight')
            plt.close()
            
            self.logger.info(f"权益曲线已保存: {os.path.join(self.output_dir, filename)}")
        except Exception as e:
            self.logger.error(f"绘制权益曲线失败: {str(e)}", exc_info=True)
    
    def plot_drawdown(self, backtest_data: pd.DataFrame, title: str = "回撤分析", 
                     filename: str = "drawdown.png") -> None:
        """
        绘制回撤曲线
        
        参数:
            backtest_data: 回测数据
            title: 图表标题
            filename: 输出文件名
        """
        try:
            plt.figure(figsize=(12, 6))
            
            # 计算回撤
            if 'cum_returns' not in backtest_data.columns:
                backtest_data['cum_returns'] = (1 + backtest_data['returns']).cumprod()
                
            cummax = backtest_data['cum_returns'].cummax()
            drawdown = (cummax - backtest_data['cum_returns']) / cummax
            
            # 绘制回撤曲线
            plt.fill_between(backtest_data.index, 0, drawdown, alpha=0.3, color='r')
            plt.plot(backtest_data.index, drawdown, color='r', alpha=0.8, linewidth=1)
            
            plt.title(title, fontsize=14)
            plt.xlabel('日期', fontsize=12)
            plt.ylabel('回撤', fontsize=12)
            plt.grid(True)
            
            # 保存图表
            plt.savefig(os.path.join(self.output_dir, filename), dpi=300, bbox_inches='tight')
            plt.close()
            
            self.logger.info(f"回撤曲线已保存: {os.path.join(self.output_dir, filename)}")
        except Exception as e:
            self.logger.error(f"绘制回撤曲线失败: {str(e)}", exc_info=True)
    
    def plot_monthly_returns(self, backtest_data: pd.DataFrame, title: str = "月度收益热图", 
                            filename: str = "monthly_returns.png") -> None:
        """
        绘制月度收益热图
        
        参数:
            backtest_data: 回测数据
            title: 图表标题
            filename: 输出文件名
        """
        try:
            # 确保数据有日期索引
            if not isinstance(backtest_data.index, pd.DatetimeIndex):
                self.logger.warning("数据没有日期索引，无法绘制月度收益热图")
                return
                
            # 计算月度收益
            monthly_returns = backtest_data['returns'].resample('M').apply(lambda x: (1 + x).prod() - 1)
            
            # 转换为年月格式的表格
            monthly_returns_table = monthly_returns.unstack(level=0)
            
            plt.figure(figsize=(12, 8))
            
            # 绘制热图
            import seaborn as sns
            monthly_pivot = pd.pivot_table(
                monthly_returns.reset_index(),
                values='returns',
                index=monthly_returns.index.month,
                columns=monthly_returns.index.year
            )
            
            # 设置月份标签
            month_labels = ['一月', '二月', '三月', '四月', '五月', '六月', 
                           '七月', '八月', '九月', '十月', '十一月', '十二月']
            
            # 绘制热图
            ax = sns.heatmap(monthly_pivot, annot=True, fmt=".2%", cmap="RdYlGn",
                           linewidths=0.5, center=0, vmin=-0.1, vmax=0.1)
            
            # 设置标签
            ax.set_yticklabels(month_labels)
            plt.title(title, fontsize=14)
            plt.xlabel('年份', fontsize=12)
            plt.ylabel('月份', fontsize=12)
            
            # 保存图表
            plt.savefig(os.path.join(self.output_dir, filename), dpi=300, bbox_inches='tight')
            plt.close()
            
            self.logger.info(f"月度收益热图已保存: {os.path.join(self.output_dir, filename)}")
        except Exception as e:
            self.logger.error(f"绘制月度收益热图失败: {str(e)}", exc_info=True)
    
    def plot_rolling_sharpe(self, backtest_data: pd.DataFrame, window: int = 60, 
                           title: str = "滚动夏普比率", filename: str = "rolling_sharpe.png") -> None:
        """
        绘制滚动夏普比率
        
        参数:
            backtest_data: 回测数据
            window: 滚动窗口大小
            title: 图表标题
            filename: 输出文件名
        """
        try:
            plt.figure(figsize=(12, 6))
            
            # 计算滚动夏普比率
            rolling_mean = backtest_data['returns'].rolling(window=window).mean() * 252
            rolling_std = backtest_data['returns'].rolling(window=window).std() * np.sqrt(252)
            rolling_sharpe = rolling_mean / rolling_std
            
            # 绘制滚动夏普比率
            plt.plot(backtest_data.index, rolling_sharpe, linewidth=2)
            plt.axhline(y=1, color='r', linestyle='--', alpha=0.7)
            plt.axhline(y=2, color='g', linestyle='--', alpha=0.7)
            plt.axhline(y=0, color='k', linestyle='-', alpha=0.2)
            
            plt.title(title, fontsize=14)
            plt.xlabel('日期', fontsize=12)
            plt.ylabel('夏普比率', fontsize=12)
            plt.grid(True)
            
            # 保存图表
            plt.savefig(os.path.join(self.output_dir, filename), dpi=300, bbox_inches='tight')
            plt.close()
            
            self.logger.info(f"滚动夏普比率已保存: {os.path.join(self.output_dir, filename)}")
        except Exception as e:
            self.logger.error(f"绘制滚动夏普比率失败: {str(e)}", exc_info=True)
    
    def plot_returns_distribution(self, backtest_data: pd.DataFrame, benchmark_data: Optional[pd.DataFrame] = None,
                                title: str = "收益分布", filename: str = "returns_distribution.png") -> None:
        """
        绘制收益分布直方图
        
        参数:
            backtest_data: 回测数据
            benchmark_data: 基准数据
            title: 图表标题
            filename: 输出文件名
        """
        try:
            plt.figure(figsize=(12, 6))
            
            # 绘制策略收益分布
            plt.hist(backtest_data['returns'], bins=50, alpha=0.5, label='策略', density=True)
            
            # 绘制基准收益分布
            if benchmark_data is not None:
                plt.hist(benchmark_data['returns'], bins=50, alpha=0.5, label='基准', density=True)
            
            # 绘制正态分布拟合曲线
            from scipy import stats
            mu, sigma = stats.norm.fit(backtest_data['returns'])
            x = np.linspace(min(backtest_data['returns']), max(backtest_data['returns']), 100)
            plt.plot(x, stats.norm.pdf(x, mu, sigma), 'k', linewidth=2, label='正态分布拟合')
            
            plt.title(title, fontsize=14)
            plt.xlabel('日收益率', fontsize=12)
            plt.ylabel('频率', fontsize=12)
            plt.legend()
            plt.grid(True)
            
            # 保存图表
            plt.savefig(os.path.join(self.output_dir, filename), dpi=300, bbox_inches='tight')
            plt.close()
            
            self.logger.info(f"收益分布直方图已保存: {os.path.join(self.output_dir, filename)}")
        except Exception as e:
            self.logger.error(f"绘制收益分布直方图失败: {str(e)}", exc_info=True)
    
    def plot_metrics_radar(self, metrics: Dict[str, float], benchmark_metrics: Optional[Dict[str, float]] = None,
                          title: str = "绩效指标雷达图", filename: str = "metrics_radar.png") -> None:
        """
        绘制绩效指标雷达图
        
        参数:
            metrics: 策略绩效指标
            benchmark_metrics: 基准绩效指标
            title: 图表标题
            filename: 输出文件名
        """
        try:
            # 选择要展示的指标
            selected_metrics = [
                "年化收益率", "夏普比率", "索提诺比率", "卡玛比率", 
                "胜率", "最大回撤", "波动率"
            ]
            
            # 准备数据
            metrics_values = []
            for metric in selected_metrics:
                if metric in metrics:
                    # 对于回撤和波动率，使用其倒数（越小越好）
                    if metric in ["最大回撤", "波动率"]:
                        value = 1 / (metrics[metric] + 0.0001)  # 避免除以零
                    else:
                        value = metrics[metric]
                    metrics_values.append(value)
                else:
                    metrics_values.append(0)
            
            # 如果有基准指标，也准备基准数据
            if benchmark_metrics:
                benchmark_values = []
                for metric in selected_metrics:
                    if metric in benchmark_metrics:
                        if metric in ["最大回撤", "波动率"]:
                            value = 1 / (benchmark_metrics[metric] + 0.0001)
                        else:
                            value = benchmark_metrics[metric]
                        benchmark_values.append(value)
                    else:
                        benchmark_values.append(0)
            
            # 设置雷达图
            angles = np.linspace(0, 2*np.pi, len(selected_metrics), endpoint=False).tolist()
            angles += angles[:1]  # 闭合图形
            
            metrics_values += metrics_values[:1]  # 闭合数据
            if benchmark_metrics:
                benchmark_values += benchmark_values[:1]
            
            fig, ax = plt.subplots(figsize=(10, 8), subplot_kw=dict(polar=True))
            
            # 绘制策略指标
            ax.plot(angles, metrics_values, 'o-', linewidth=2, label='策略')
            ax.fill(angles, metrics_values, alpha=0.25)
            
            # 绘制基准指标
            if benchmark_metrics:
                ax.plot(angles, benchmark_values, 'o-', linewidth=2, label='基准')
                ax.fill(angles, benchmark_values, alpha=0.1)
            
            # 设置标签
            ax.set_thetagrids(np.degrees(angles[:-1]), selected_metrics)
            
            plt.title(title, fontsize=14)
            plt.legend(loc='upper right')
            
            # 保存图表
            plt.savefig(os.path.join(self.output_dir, filename), dpi=300, bbox_inches='tight')
            plt.close()
            
            self.logger.info(f"绩效指标雷达图已保存: {os.path.join(self.output_dir, filename)}")
        except Exception as e:
            self.logger.error(f"绘制绩效指标雷达图失败: {str(e)}", exc_info=True)
    
    def generate_report(self, backtest_data: pd.DataFrame, metrics: Dict[str, float], 
                       benchmark_data: Optional[pd.DataFrame] = None, 
                       benchmark_metrics: Optional[Dict[str, float]] = None,
                       strategy_name: str = "未命名策略") -> None:
        """
        生成完整的回测报告
        
        参数:
            backtest_data: 回测数据
            metrics: 策略绩效指标
            benchmark_data: 基准数据
            benchmark_metrics: 基准绩效指标
            strategy_name: 策略名称
        """
        try:
            # 绘制各种图表
            self.plot_equity_curve(backtest_data, benchmark_data, 
                                 title=f"{strategy_name} - 权益曲线")
            
            self.plot_drawdown(backtest_data, 
                              title=f"{strategy_name} - 回撤分析")
            
            self.plot_monthly_returns(backtest_data, 
                                    title=f"{strategy_name} - 月度收益热图")
            
            self.plot_rolling_sharpe(backtest_data, 
                                   title=f"{strategy_name} - 滚动夏普比率")
            
            self.plot_returns_distribution(backtest_data, benchmark_data, 
                                         title=f"{strategy_name} - 收益分布")
            
            self.plot_metrics_radar(metrics, benchmark_metrics, 
                                  title=f"{strategy_name} - 绩效指标雷达图")
            
            # 生成HTML报告
            self._generate_html_report(backtest_data, metrics, benchmark_data, 
                                     benchmark_metrics, strategy_name)
            
            self.logger.info(f"回测报告已生成: {self.output_dir}")
        except Exception as e:
            self.logger.error(f"生成回测报告失败: {str(e)}", exc_info=True)
    
    def _generate_html_report(self, backtest_data: pd.DataFrame, metrics: Dict[str, float], 
                            benchmark_data: Optional[pd.DataFrame] = None, 
                            benchmark_metrics: Optional[Dict[str, float]] = None,
                            strategy_name: str = "未命名策略") -> None:
        """
        生成HTML格式的回测报告
        
        参数:
            backtest_data: 回测数据
            metrics: 策略绩效指标
            benchmark_data: 基准数据
            benchmark_metrics: 基准绩效指标
            strategy_name: 策略名称
        """
        try:
            # 准备HTML模板
            html_template = """
            <!DOCTYPE html>
            <html>
            <head>
                <meta charset="UTF-8">
                <meta name="viewport" content="width=device-width, initial-scale=1.0">
                <title>{strategy_name} - 回测报告</title>
                <style>
                    body {{ font-family: Arial, sans-serif; margin: 0; padding: 20px; }}
                    h1, h2, h3 {{ color: #333; }}
                    .container {{ max-width: 1200px; margin: 0 auto; }}
                    .metrics-table {{ width: 100%; border-collapse: collapse; margin-bottom: 20px; }}
                    .metrics-table th, .metrics-table td {{ border: 1px solid #ddd; padding: 8px; text-align: left; }}
                    .metrics-table th {{ background-color: #f2f2f2; }}
                    .chart-container {{ margin-bottom: 30px; }}
                    .chart {{ width: 100%; max-width: 800px; margin: 0 auto; }}
                    .chart img {{ width: 100%; height: auto; border: 1px solid #ddd; }}
                </style>
            </head>
            <body>
                <div class="container">
                    <h1>{strategy_name} - 回测报告</h1>
                    <p>回测期间: {start_date} 至 {end_date}</p>
                    
                    <h2>绩效指标</h2>
                    <table class="metrics-table">
                        <tr>
                            <th>指标</th>
                            <th>策略</th>
                            {benchmark_header}
                        </tr>
                        {metrics_rows}
                    </table>
                    
                    <h2>图表分析</h2>
                    
                    <div class="chart-container">
                        <h3>权益曲线</h3>
                        <div class="chart">
                            <img src="equity_curve.png" alt="权益曲线">
                        </div>
                    </div>
                    
                    <div class="chart-container">
                        <h3>回撤分析</h3>
                        <div class="chart">
                            <img src="drawdown.png" alt="回撤分析">
                        </div>
                    </div>
                    
                    <div class="chart-container">
                        <h3>月度收益热图</h3>
                        <div class="chart">
                            <img src="monthly_returns.png" alt="月度收益热图">
                        </div>
                    </div>
                    
                    <div class="chart-container">
                        <h3>滚动夏普比率</h3>
                        <div class="chart">
                            <img src="rolling_sharpe.png" alt="滚动夏普比率">
                        </div>
                    </div>
                    
                    <div class="chart-container">
                        <h3>收益分布</h3>
                        <div class="chart">
                            <img src="returns_distribution.png" alt="收益分布">
                        </div>
                    </div>
                    
                    <div class="chart-container">
                        <h3>绩效指标雷达图</h3>
                        <div class="chart">
                            <img src="metrics_radar.png" alt="绩效指标雷达图">
                        </div>
                    </div>
                </div>
            </body>
            </html>
            """
            
            # 准备指标行
            metrics_rows = ""
            for key, value in metrics.items():
                if isinstance(value, float):
                    benchmark_value = f"<td>{benchmark_metrics.get(key, 0):.4f}</td>" if benchmark_metrics else ""
                    metrics_rows += f"<tr><td>{key}</td><td>{value:.4f}</td>{benchmark_value}</tr>\n"
                else:
                    benchmark_value = f"<td>{benchmark_metrics.get(key, '')}</td>" if benchmark_metrics else ""
                    metrics_rows += f"<tr><td>{key}</td><td>{value}</td>{benchmark_value}</tr>\n"
            
            # 填充模板
            html_content = html_template.format(
                strategy_name=strategy_name,
                start_date=backtest_data.index[0].strftime('%Y-%m-%d') if isinstance(backtest_data.index, pd.DatetimeIndex) else "N/A",
                end_date=backtest_data.index[-1].strftime('%Y-%m-%d') if isinstance(backtest_data.index, pd.DatetimeIndex) else "N/A",
                benchmark_header="<th>基准</th>" if benchmark_metrics else "",
                metrics_rows=metrics_rows
            )
            
            # 保存HTML报告
            with open(os.path.join(self.output_dir, "report.html"), "w", encoding="utf-8") as f:
                f.write(html_content)
                
            self.logger.info(f"HTML报告已生成: {os.path.join(self.output_dir, 'report.html')}")
        except Exception as e:
            self.logger.error(f"生成HTML报告失败: {str(e)}", exc_info=True)
